from preprocessing import parse_annotation, BatchGenerator
import numpy as np
from frontend import YOLO
from matplotlib import pyplot as plt
import sys

RBC_DATA_DIR = '../../../../../large_data/CV4/_many_files/RBC_datasets'
MODEL_PATH = '../../../../../large_data/model/yolo_v2/keras_yolo2-pre-trained-weights/tiny_yolo_backend.h5'

config = dict()
config["train"] = dict()
config["model"] = dict()
config["valid"] = dict()
config["valid"]["valid_times"] = 1
# config["train"]["train_annot_folder"]="RBC_datasets/Annotations/"
config["train"]["train_annot_folder"] = RBC_DATA_DIR + "/Annotations/"
# config["train"]["train_image_folder"]="RBC_datasets/JPEGImages/"
config["train"]["train_image_folder"] = RBC_DATA_DIR + "/JPEGImages/"
config["train"]["train_times"] = 1
config["train"]["batch_size"] = 16
config["train"]["learning_rate"] = 1e-4
config["train"]["nb_epoch"] = 2  # ori 50
config["train"]["warmup_batches"] = 250
config["train"]["object_scale"] = 5.0
config["train"]["no_object_scale"] = 1.0
config["train"]["coord_scale"] = 1.0
config["train"]["class_scale"] = 1.0
config["model"]["labels"] = ["RBC"]
config["model"]["input_size"] = 416
config["model"]["max_box_per_img"] = 10
config["model"]["anchors"] = [0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828]
config["model"]["architecture"] = "Tiny Yolo"
config["model"]["grid_w"] = 13
config["model"]["grid_h"] = 13

# weights_path = "models/tiny_yolo.h5"
weights_path = MODEL_PATH

imgs,_=parse_annotation(config["train"]["train_annot_folder"],
                        config["train"]["train_image_folder"],
                        config["model"]["labels"]
                       )

# print(imgs[:2])
# print(_)
# sys.exit(0)

split_idx=int(0.8*len(imgs))
np.random.shuffle(imgs)
train_imgs=imgs[:split_idx]
valid_imgs=imgs[split_idx:]
print(train_imgs[0])

generator_config={
        "IMAGE_H":config["model"]["input_size"],
        "IMAGE_W":config["model"]["input_size"],
        "GRID_H":config["model"]["grid_h"],
        "GRID_W":config["model"]["grid_w"],
        "BOX":5,
        "LABELS":config["model"]["labels"],
        "CLASS":len(config["model"]["labels"]),
        "ANCHORS":config["model"]["anchors"],
        "BATCH_SIZE":config["train"]["batch_size"],
        "TRUE_BOX_BUFFER":config["model"]["max_box_per_img"]
        }

generator=BatchGenerator(train_imgs,generator_config,norm=None)
[x,b],y=generator.__getitem__(0)
print('x', x.shape)
print('b', b.shape)
print('y', y.shape)
print('x[0]', x[0].shape)
print('b[0]', b[0].shape)
print('y[0]', y[0].shape)
plt.imshow(x[0])
plt.show()
# sys.exit(0)

yolo=YOLO(architecture=config["model"]["architecture"],
            input_size=config["model"]["input_size"],
            labels=config["model"]["labels"],
            max_box_per_img=config["model"]["max_box_per_img"],
            anchors=config["model"]["anchors"])

yolo.train(
    train_imgs,
    valid_imgs,
    config["train"]["train_times"],
    config["valid"]["valid_times"],
    config["train"]["nb_epoch"],
    config["train"]["learning_rate"],
    config["train"]["batch_size"],
    config["train"]["warmup_batches"],
    config["train"]["object_scale"],
    config["train"]["no_object_scale"],
    config["train"]["coord_scale"],
    config["train"]["class_scale"],
    # train=False,  # ori
    train=True,
)
# yolo.load_weights(weights_path)  # ori

[test_x,test_b],test_y=yolo.compute_loss(train_imgs)
netout=yolo.model.predict([np.expand_dims(test_x,0),np.expand_dims(test_b,0)])[0]

from utils import decode_netout,draw_boxes
boxes=decode_netout(netout,0.51,0.2,config["model"]["anchors"],len(config["model"]["labels"]))
test_x_copy=test_x.copy()
img=draw_boxes(boxes,test_x_copy,config["model"]["labels"])
plt.imshow(img)
plt.show()
